#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/sort.h>
#include <thrust/execution_policy.h>
#include <thread>
#include <thrust/reduce.h>
#include <thrust/functional.h>
#include <thrust/execution_policy.h>
#include "cub/block/block_radix_sort.cuh"
struct iAS {
    int idx;
    int score;
    __host__ __device__
    bool operator < (const iAS &x) const {
        if(score != x.score) return score < x.score;
        return idx > x.idx;
    }
};

struct thrust_cmp
{
    // __host__ __device__
    // bool operator()(const iAS &a, const iAS &b) {
    //     // return a < b;
    //     if(a.score != b.score) return a.score > b.score;
    //     return a.idx < b.idx;
    // }

    __host__ __device__
    iAS operator()(const iAS &a, const iAS &b) {
        // return a < b;
        if(a < b) return b;
        return a;
    }
};


// CUDA内核函数，使用thrust::stable_sort排序
__global__ void sort_kernel(int* d_idx, float* d_score, int n) {
    thrust::stable_sort_by_key(thrust::device, d_score, 
        d_score + n, d_idx, thrust::greater<float>());
}

void cpu_sort(int* d_idx, float* d_score, int n) {
    // thrust::stable_sort_by_key(thrust::device, d_score, 
    //     d_score + n, d_idx, thrust::greater<float>());
}


int main() {
    int n = 7853052;  // 数据大小
    // thrust::host_vector<int> h_data(n);  // 主机端数据
    // thrust::device_vector<int> d_data;  // 设备端数据

    cudaDeviceProp device_props;
    cudaGetDeviceProperties(&device_props, 0);

    cudaSetDevice(0);


    long long *h_vector_idx = new long long[n];
    long long *h_vector_score = new long long[n];
    long long *d_vector_idx = nullptr; cudaMalloc(&d_vector_idx, sizeof(long long) * n);
    long long *d_vector_score = nullptr; cudaMalloc(&d_vector_score, sizeof(long long) * n);

    // 初始化数据
    for (int i = 0; i < n; i++) {
        h_vector_idx[i] = i;
        h_vector_score[i] = i;
    }
    
    // 将数据从主机端传输到设备端
    // d_data = h_data;

    cudaMemcpy(d_vector_idx, h_vector_idx, sizeof(long long) * n, cudaMemcpyHostToDevice);
    cudaMemcpy(d_vector_score, h_vector_score, sizeof(long long) * n, cudaMemcpyHostToDevice);
    clock_t time1,time2;

    // 在设备端使用thrust::stable_sort进行排序

    int group = 2;
    time1 = clock();
    // std::thread threads[group];
    // for(int i = 0; i < group; i++) {
    //     int group_size = (n + group - 1) / group;
    //     int bound = min(n - i * group_size, group_size);
        // std::cout << group_size << " " << bound << " " << i * group_size << "\n";
        // threads[i] = std::thread(cpu_sort, d_vector_idx + i * group_size, d_vector_score + i * group_size, bound);
        // threads[i] = std::thread(thrust::stable_sort_by_key, d_vector_score + i * group_size, 
        //     d_vector_score + i * group_size + bound, d_vector_idx + i * group_size, thrust::greater<float>());
        // sort_kernel<<<1, 1>>>(d_vector_idx + i * group_size, d_vector_score + i * group_size, bound);  // 设置核函数的块和线程数，这里使用256个线程块，每个块大小为1。
        // thrust::stable_sort_by_key(thrust::device, d_vector_score + i * group_size, 
        //     d_vector_score + i * group_size + bound, d_vector_idx + i * group_size, thrust::greater<float>());
    // }
    
    // for(int i = 0; i < group; i++) {
    //     threads[i].join();
    // }

    // sort_kernel<<<1, 1>>>(d_vector_idx, d_vector_score, n);  // 设置核函数的块和线程数，这里使用256个线程块，每个块大小为1。
    // cudaDeviceSynchronize();
    // for(int k = 0; k < 100; k++) {
        // thrust::device::sort(thrust::device, d_vector_score, d_vector_score + n);
    // }
    // test1.join();
    typedef cub::BlockRadixSort<int, 128, 4> BlockRadixSort;
    BlockRadixSort(h_vector_idx);
    time2 = clock();
    std::cout<<(double)(time2-time1)/CLOCKS_PER_SEC<<std::endl;
    
    
    // 将数据从设备端传输回主机端并打印结果
    // cudaMemcpy(h_vector_idx, d_vector_idx, sizeof(int) * n, cudaMemcpyDeviceToHost);
    // cudaMemcpy(h_vector_score, d_vector_score, sizeof(int) * n, cudaMemcpyDeviceToHost);
    
    // for (int i = 0; i < 5; i++) {
    //     std::cout << h_vector_idx[i] << " " << h_vector_score[i] << "\n";
    // }
    // thrust::copy(d_vector, d_vector + n, std::ostream_iterator<int>(std::cout, " "));
    // std::cout << std::endl;
    
    return 0;
}